'''
Created on Aug 6, 2014

@author: Siyuan Qi
'''

# Dependencies
import pymc.distributions

# Local imports
from entity import Entity

# ====================== Classes of different kinds of nodes ======================
class Terminal:
	def __init__(self, node, distribution, cameraHeight):
		""" Creates a terminal node.

		Creates a terminal node in the grammar tree.

		Args:
			node: The corresponding XML tree node extracted from <Symbols>, which contains the
				rotation center of the object, if the node is an entity, and the render color of 
				the object.
			distribution: The corresponding XML tree node extracted from <Distributions>, which contains
				the parameters of the distributions of position, size and orientation
			cameraHeight: The camera height in the scene.  Since the coordinates in the proposed 
				cuboids are based on the coordinate system using the camera as the origin, the mean 
				position of 'z' value of each object will be deducted by cameraHeight.
		"""
		self.symbol = node.text
		self.anchor = [float(a) for a in node.attrib['anchor'].split()]
		self.isEntity = float(node.attrib['isEntity'])
		self.isTerminal = True

		self.parent = None
		self.children = None

		self.distribution = distribution

		if self.isEntity == 1:
			self.color = [int(float(c)) for c in node.attrib['color'].split()]

		for i in range(0, len(self.distribution.zMu)):
			self.distribution.zMu[i] -= cameraHeight

	def sample(self, entity):
		""" Sample the children entities of an entity of this grammar node.

		Sample the children entites of an entity of this grammar node, given that the entity is sampled.
		Since this is an terminal entity, the children will be assigned with None.

		Args:
			entity: An entity of this grammar node.  The entity is already sampled with attributes --
				the position, size and orientation are already sampled.

		Returns:
			The entity corresponding to this grammar node with children entities sampled.

		"""
		entity.children = None
		return entity


class SetProduction:
	def __init__(self, node, children, childrenProb, distribution, cameraHeight):	
		""" Creates a non-terminal node.

		Creates a non-terminal node in the grammar tree.

		Args:
			node: The corresponding XML tree node extracted from <Symbols>, which contains the
				rotation center of the object, if the node is an entity, and the render color of 
				the object.
			children: The children nodes of the set node.  Notice that the children nodes are the nodes
				in the grammar tree, instead of the XML tree or the parse tree.
			childrenProb: The probability distribution of the number of each child.
			distribution: The corresponding XML tree node extracted from <Distributions>, which contains
				the parameters of the distributions of position, size and orientation
			cameraHeight: The camera height in the scene.  Since the coordinates in the proposed 
				cuboids are based on the coordinate system using the camera as the origin, the mean 
				position of 'z' value of each object will be deducted by cameraHeight.
		"""
		self.symbol = node.text
		self.anchor = [float(a) for a in node.attrib['anchor'].split()]
		self.isEntity = float(node.attrib['isEntity'])
		self.isTerminal = False

		self.parent = None
		self.children = children
		self.childrenProb = childrenProb

		self.distribution = distribution

		if self.isEntity == 1:
			self.color = [int(float(c)) for c in node.attrib['color'].split()]

		for i in range(0, len(self.distribution.zMu)):
			self.distribution.zMu[i] -= cameraHeight

	def sample(self, entity):
		""" Sample the children entities of an entity of this grammar node.

		Sample the children entites of an entity of this grammar node, given that the entity is sampled.
		Notice that this function calls "sampleAttribute()" method of Entity, and it samples the whole
		tree of children instead of a single layer.

		Args:
			entity: An entity of this grammar node.  The entity is already sampled with attributes --
				the position, size and orientation are already sampled.

		Returns:
			The entity corresponding to this grammar node with children entities sampled.
		"""
		# For each child, sample the quantity according to the defined probability
		childrenNodes = []
		childrenCount = {}
		for child in self.children:
			childrenCount[child.symbol] = 0
			for i in range(0, pymc.distributions.rcategorical(self.childrenProb[child.symbol])):
				childrenNodes.append(child)
				childrenCount[child.symbol] += 1

		# Sample the entities of children nodes
		childrenEntities = []
		for childNode in childrenNodes:
			childEntity = Entity(childNode, entity)
			childrenEntities.append(childEntity)

			childEntity.sampleAttribute()
			#while not (childEntity.compatibleWithParent() and mutualExclusive(childrenEntities)):
			while not True: #mutualExclusive(childrenEntities): #TODO
				childEntity.sampleAttribute()

			childNode.sample(childEntity)

		entity.children = childrenEntities
		entity.childrenCount = childrenCount
		return entity







